#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
@File        : microbiome_bracken_merge.py
@Author      : Bing Liang
@Email       : believer19940901@gmail.com
@Date        : 2025/10/10
@Description : 合并多个样本的物种丰度（RPM）
               - 支持最小RPM过滤
               - 支持最小出现样本数过滤
               - 自动剔除全为零的样本列
"""

from argparse import ArgumentParser, Namespace
from pathlib import Path
import pandas as pd
import sys


def main(args: Namespace):
    """主函数：合并所有样本的物种丰度（RPM）"""
    min_value = max(args.min_value, 0)
    min_count = max(args.min_count, 0)

    # 解析样本名与文件路径
    samples = [s.strip() for s in args.samples.split(",") if s.strip()]
    infiles = [Path(f.strip()).resolve() for f in args.input.split(",") if f.strip()]

    # 参数检查
    if len(samples) != len(infiles):
        sys.exit(f"样本数量 ({len(samples)}) 与输入文件数量 ({len(infiles)}) 不一致。")

    # 创建输出目录
    outfile = Path(args.output).resolve()
    outfile.parent.mkdir(parents=True, exist_ok=True)

    all_data = {}
    skipped = []

    # 逐个读取输入文件
    for sample, infile in zip(samples, infiles):
        if not infile.exists():
            print(f"文件不存在，跳过样本 {sample}: {infile}")
            skipped.append(sample)
            continue

        try:
            df = pd.read_csv(infile, sep="\t", usecols=["species", "RPM"])
        except ValueError as e:
            sys.exit(f"文件 {infile} 缺少必要列（species, RPM）。错误详情: {e}")
        except Exception as e:
            sys.exit(f"读取文件 {infile} 失败: {e}")

        # 过滤低丰度物种
        df = df[df["RPM"] >= min_value]
        all_data[sample] = df.set_index("species")["RPM"]

    if not all_data:
        sys.exit("没有可用的输入文件，无法生成输出结果。")

    # 合并所有样本（按物种对齐）
    merged_df = pd.concat(all_data, axis=1)
    merged_df.columns = merged_df.columns.get_level_values(0)
    merged_df.fillna(0, inplace=True)
    merged_df.reset_index(inplace=True)

    # 保留至少出现在 min_count 个样本中的物种
    presence_mask = (merged_df.iloc[:, 1:] > 0).sum(axis=1) >= min_count
    filtered_df = merged_df.loc[presence_mask, :]

    # 检查并记录全为零的样本列
    zero_sum_cols = [col for col in filtered_df.columns[1:] if filtered_df[col].sum() == 0]

    # 删除全为零列
    filtered_df = filtered_df.loc[:, ["species"] + [
        col for col in filtered_df.columns[1:] if col not in zero_sum_cols
    ]]

    # 输出结果
    filtered_df.to_csv(outfile, sep="\t", index=False)

    print(f"输出文件已生成: {outfile}")
    print(f"合并样本数: {len(samples)}（跳过 {len(skipped)} 个）")
    print(f"输出物种数: {filtered_df.shape[0]}")

    if skipped:
        print(f"跳过的样本（文件不存在）: {', '.join(skipped)}")
    if zero_sum_cols:
        print(f"以被过滤掉的样本（全为 0）: {', '.join(zero_sum_cols)}")
    else:
        print("无样本在过滤后被移除。")


if __name__ == "__main__":
    parser = ArgumentParser(description="合并多个样本的物种丰度（RPM）")
    parser.add_argument("-s", "--samples", required=True, help="样本名（逗号分隔）")
    parser.add_argument("-i", "--input", required=True, help="输入文件路径（逗号分隔）")
    parser.add_argument("-o", "--output", required=True, help="输出文件路径（TSV格式）")
    parser.add_argument("-m", "--min_value", type=float, default=0, help="RPM过滤阈值（默认: 0）")
    parser.add_argument("-r", "--min_count", type=int, default=3, help="物种至少出现的样本数（默认: 3）")

    main(parser.parse_args())
